#include "DaeIO.hpp"
#include <Xml/RapidXML.hpp>
#include <Utility/FileTools.hpp>
#include "../Material.hpp"
#include "../../ResourceManager.hpp"
#include "../../../Context/Context.hpp"
#include "../../../Graphics/GeometryHelper.hpp"

namespace zzz{
bool DaeIO::load_mat_=true;
bool DaeIO::save_mat_=true;
string DaeIO::dae_path_;

string DaeIO::ToId(const string &str)
{
  if (str[0]=='#') return string(str.begin()+1,str.end());
  return str;
}

string DaeIO::ToSource(const string &str)
{
  if (str[0]=='#') return str;
  string ret="#";
  return ret + str;
}

RapidXMLNode DaeIO::FindNodeById(RapidXMLNode &father, const string &name, string &id)
{
  string myid = ToId(id);
  for (RapidXMLNode node=father.GetFirstNode();node.IsValid();node=node.GetNextSibling()) {
    if (name == node.GetName() && myid == node.GetAttribute("id")) return node;
  }
  return RapidXMLNode(NULL, NULL);
}

const int SEM_VERTEX=0;
const int SEM_TEXCOORD=1;
const int SEM_NORMAL=2;
const int SEM_COLOR=3;
string dae_path;

void DaeIO::ReadMaterial(TriMesh &mesh, DaeData &dae, string &url)
{
  RapidXMLNode matNode=FindNodeById(dae.libMaterialsNode,"material",url);
  string effUrl=matNode.GetFirstNode().GetAttribute("url");
  RapidXMLNode effNode(FindNodeById(dae.libEffectsNode,"effect",effUrl));
  RapidXMLNode phongNode=effNode.GetFirstNode().GetFirstNode().GetFirstNode();
  Material *mat=new Material;
  if (phongNode.HasNode("emission")) {
    RapidXMLNode node=phongNode.GetNode("emission").GetFirstNode();
    if (strcmp(node.GetName(),"color")==0) {
      string data=node.GetText();
      istringstream iss(data);
      iss>>mat->emission_;
      mat->SetFlag(MAT_EMI);
    }
  }
  if (phongNode.HasNode("ambient")) {
    RapidXMLNode node=phongNode.GetNode("ambient").GetFirstNode();
    if (strcmp(node.GetName(),"color")==0) {
      string data=node.GetText();
      istringstream iss(data);
      iss>>mat->ambient_;
      mat->SetFlag(MAT_AMB);
    } else if (strcmp(node.GetName(),"texture")==0) {
      //read texture
      string tex_id=node.GetAttribute("texture");
      string filename=FindNodeById(dae.libImagesNode,"image",tex_id).GetFirstNode().GetText();
      filename=PathFile(dae_path,filename);
      if (GLExists()) {
        Image3uc img(filename.c_str());
        mat->ambientTex_.ImageToTexture(img);
        mat->SetFlag(MAT_AMBTEX);
      } else {
        mat->ambientTexName_=filename;
      }
    }
  }
  ZRM->Add(mat,url);
}

void DaeIO::ReadGeometry(TriMesh &mesh, DaeData &dae, string &url)
{
  RapidXMLNode geoNode(FindNodeById(dae.libGeometriesNode,"geometry",url));
  RapidXMLNode meshNode=geoNode.GetNode("mesh");

  map<string,int> bases;  //input base number

  for (RapidXMLNode triNode=meshNode.GetFirstNode("triangles");triNode.IsValid();triNode.GotoNextSibling("triangles")) {
    int inputn=0;  //input number
    int inputt[10];  //input type
    int inputb[10];
    //load source and set offsets
    for (zuint i=0; i<triNode.NodeNumber(); i++) {
      RapidXMLNode node=triNode.GetNode(i);
      if (strcmp(node.GetName(),"input")!=0) continue;
      int offset=FromString<int>(node.GetAttribute("offset"));
      if (offset>=inputn) inputn=offset+1;
      string semantic=node.GetAttribute("semantic");
      //set offsets
      int base;
      if (semantic=="VERTEX") {
        inputt[offset]=SEM_VERTEX;
        base=mesh.pos_.size();
      } else if (semantic=="NORMAL") {
        inputt[offset]=SEM_NORMAL;
        base=mesh.nor_.size();
      } else if (semantic=="TEXCOORD") {
        inputt[offset]=SEM_TEXCOORD;
        base=mesh.tex_.size();
      } else if (semantic=="COLOR") {
        inputt[offset]=SEM_COLOR;
        base=mesh.color_.size();
      } else {
        ZLOGF<<"Unknown semantic: "<<semantic<<endl;
        continue;
      }

      //load source
      string source=node.GetAttribute("source");
      if (bases.find(ToId(source))!=bases.end()) {   //already loaded
        inputb[offset]=bases[ToId(source)]; 
        continue;
      }
      bases[ToId(source)]=base;
      inputb[offset]=base;
      if (inputt[offset]==SEM_VERTEX) {
        RapidXMLNode vertexNode=FindNodeById(meshNode,"vertices",source);
        string possource=vertexNode.GetFirstNode().GetAttribute("source");
        RapidXMLNode sourceNode=FindNodeById(meshNode,"source",possource);
        RapidXMLNode arrayNode=sourceNode.GetNode("float_array");
        int arrayn=FromString<int>(arrayNode.GetAttribute("count"))/3;
        string data=arrayNode.GetText();
        istringstream iss(data);
        Vector3f v;
        for (int j=0; j<arrayn; j++) {
          iss>>v;
          if (iss.fail()) {
            ZLOGE<<"Error in DaeIO::Load. \"count\" in <float_array> of "<<possource<<" does not match the actual data size."<<endl;
            break;
          }
          mesh.pos_.push_back(v);
        }
      } else if (inputt[offset]==SEM_NORMAL) {
        RapidXMLNode sourceNode=FindNodeById(meshNode,"source",source);
        RapidXMLNode arrayNode=sourceNode.GetNode("float_array");
        int arrayn=FromString<int>(arrayNode.GetAttribute("count"))/3;
        string data=arrayNode.GetText();
        istringstream iss(data);
        Vector3f v;
        for (int j=0; j<arrayn; j++) {
          iss>>v;
          if (iss.fail()) {
            ZLOGE<<"Error in DaeIO::Load. \"count\" in <float_array> of "<<source<<" does not match the actual data size."<<endl;
            break;
          }
          mesh.nor_.push_back(v);
        }
      } else if (inputt[offset]==SEM_TEXCOORD) {
        RapidXMLNode sourceNode=FindNodeById(meshNode,"source",source);
        RapidXMLNode arrayNode=sourceNode.GetNode("float_array");
        int arrayn=FromString<int>(arrayNode.GetAttribute("count"))/2;
        string data=arrayNode.GetText();
        istringstream iss(data);
        Vector3f v(0);
        for (int j=0; j<arrayn; j++) {
          iss>>v[0]>>v[1];
          if (iss.fail()) {
            ZLOGE<<"Error in DaeIO::Load. \"count\" in <float_array> of "<<source<<" does not match the actual data size."<<endl;
            break;
          }
          mesh.tex_.push_back(v);
        }
      } else if (inputt[offset]==SEM_COLOR) {
        RapidXMLNode sourceNode=FindNodeById(meshNode,"source",source);
        RapidXMLNode arrayNode=sourceNode.GetNode("float_array");
        int arrayn=FromString<int>(arrayNode.GetAttribute("count"))/3;
        string data=arrayNode.GetText();
        istringstream iss(data);
        Vector3f v;
        for (int j=0; j<arrayn; j++) {
          iss>>v;
          if (iss.fail()) {
            ZLOGE<<"Error in DaeIO::Load. \"count\" in <float_array> of "<<source<<" does not match the actual data size."<<endl;
            break;
          }
          mesh.color_.push_back(v);
        }
      }
    }
    int trin=FromString<int>(triNode.GetAttribute("count"));
    string data=triNode.GetNode("p").GetText();
    istringstream iss(data);
    mesh.groups_.push_back(new TriMesh::MeshGroup());
    TriMesh::MeshGroup *group=mesh.groups_.back();
    int v;
    for (int i=0; i<trin; i++) {
      Vector3i facep,facet,facen,facec;
      for (int k=0; k<3; k++) {
        for (int j=0; j<inputn; j++) {
          iss>>v;
          if (iss.fail()) {
            ZLOGE<<"Error in DaeIO::Load. \"count\" in <triangles> does not match the actual data size.\n";
            goto BREAK_OUTSIDE;
          }
          if (inputt[j]==SEM_VERTEX) facep[k]=v+inputb[j];
          else if (inputt[j]==SEM_NORMAL) facen[k]=v+inputb[j];
          else if (inputt[j]==SEM_TEXCOORD) facet[k]=v+inputb[j];
          else if (inputt[j]==SEM_COLOR) facec[k]=v+inputb[j];
        }
      }
      for (int j=0; j<inputn; j++) {
        if (inputt[j]==SEM_VERTEX) group->facep_.push_back(facep);
        else if (inputt[j]==SEM_NORMAL) group->facen_.push_back(facen);
        else if (inputt[j]==SEM_TEXCOORD) group->facet_.push_back(facet);
        else if (inputt[j]==SEM_COLOR) group->facec_.push_back(facec);
      }
    }
    BREAK_OUTSIDE:
    if (load_mat_ && triNode.HasAttribute("material")) {
      string matname=triNode.GetAttribute("material");
      ReadMaterial(mesh,dae,matname);
      group->matname_=matname;
    }
  }
}

void DaeIO::ReadVisual(TriMesh &mesh, DaeData &dae, string &url)
{
  RapidXMLNode sceneNode=FindNodeById(dae.libVisualsNode,"visual_scene",url);
  for (zuint modeli=0;modeli<sceneNode.NodeNumber();modeli++)
  {
    RapidXMLNode modelNode=sceneNode.GetNode(modeli);
    string geometryUrl=modelNode.GetNode("instance_geometry").GetAttribute("url");
    ReadGeometry(mesh,dae,geometryUrl);
  }
}

bool DaeIO::Load(const string &filename, TriMesh &mesh, bool loadtex)
{
  load_mat_ = loadtex;

  RapidXML xml;
  if (!xml.LoadFile(filename)) return false;
  RapidXMLNode head=xml.GetNode("COLLADA");
  
  dae_path=GetPath(filename);
  mesh.Clear();
  DaeData dae(head.GetNode("library_images"),
        head.GetNode("library_materials"),
        head.GetNode("library_effects"),
        head.GetNode("library_geometries"),
        head.GetNode("library_visual_scenes"));
  
  RapidXMLNode sceneNode=head.GetNode("scene");
  string visualUrl=sceneNode.GetNode("instance_visual_scene").GetAttribute("url");

  ReadVisual(mesh,dae,visualUrl);

  mesh.AABB_.Reset();
  mesh.AABB_+=mesh.pos_;

  mesh.SetDataFlags();
  return true;

}
/////////////////////////////////////////////////////////////////////
void DaeIO::WriteMaterial(TriMesh &mesh, DaeData &dae, string &url)
{
  //material
  if (FindNodeById(dae.libMaterialsNode,"material",url).IsValid()) return;
  RapidXMLNode matNode=dae.libMaterialsNode.AppendNode("material");
  matNode.SetAttribute("id",url);
  matNode.SetAttribute("name",url);
  string effUrl=url+"-effect";
  matNode.AppendNode("instance_effect")<<make_pair("url",ToSource(effUrl));
  //effect
  RapidXMLNode effNode=dae.libEffectsNode.AppendNode("effect");
  effNode.SetAttribute("id",ToId(effUrl));
  effNode.SetAttribute("name",ToId(effUrl));
  RapidXMLNode techNode=effNode.AppendNode("profile_COMMON").AppendNode("technique");
  techNode.SetAttribute("sid","COMMON");
  RapidXMLNode phongNode=techNode.AppendNode("phong");
  Material *mat=ZRM->Get<Material*>(url);
  if (mat->HasFlag(MAT_EMI))
    phongNode.AppendNode("emission").AppendNode("color")<<ToString(mat->emission_);
  if (mat->HasFlag(MAT_DIF))
    phongNode.AppendNode("emission").AppendNode("color")<<ToString(mat->diffuse_);
  if (mat->HasFlag(MAT_AMB))
    phongNode.AppendNode("emission").AppendNode("color")<<ToString(mat->ambient_);
  if (mat->HasFlag(MAT_SPE))
    phongNode.AppendNode("emission").AppendNode("color")<<ToString(mat->specular_);

  if (mat->HasFlag(MAT_AMBTEX) || !mat->ambientTexName_.empty())
  {
    string imgid=ToId(url)+"-ambimg";
    phongNode.AppendNode("ambient").AppendNode("texture")<<make_pair("texture",ToId(imgid));
    RapidXMLNode imgNode=dae.libImagesNode.AppendNode("image");
    imgNode.SetAttribute("id",ToId(imgid));
    imgNode.SetAttribute("name",ToId(imgid));
    if (mat->ambientTexName_.empty())
    {
      mat->ambientTex_.TextureToFile<Vector3uc>((imgid+".png").c_str());
      imgNode.AppendNode("init_from")<<imgid+".png";
    }
    else
      imgNode.AppendNode("init_from")<<mat->ambientTexName_;
  }

  if (mat->HasFlag(MAT_DIFTEX) || !mat->diffuseTexName_.empty())
  {
    string imgid=ToId(url)+"-difimg";
    phongNode.AppendNode("diffuse").AppendNode("texture")<<make_pair("texture",ToId(imgid));
    RapidXMLNode imgNode=dae.libImagesNode.AppendNode("image");
    imgNode.SetAttribute("id",ToId(imgid));
    imgNode.SetAttribute("name",ToId(imgid));
    if (mat->diffuseTexName_.empty())
    {
      mat->diffuseTex_.TextureToFile<Vector3uc>((imgid+".png").c_str());
      imgNode.AppendNode("init_from")<<imgid+".png";

    }
    else
      imgNode.AppendNode("init_from")<<mat->diffuseTexName_;
  }

}

void DaeIO::WriteGeometry(TriMesh &mesh, DaeData &dae, string &url)
{
  RapidXMLNode geoNode=dae.libGeometriesNode.AppendNode("geometry");
  geoNode.SetAttribute("id",ToId(url));
  geoNode.SetAttribute("name",ToId(url));
  RapidXMLNode meshNode=geoNode.AppendNode("mesh");

  //for now, only output position and texcoord
  //since i don't have a file to demonstrate how normal is stored

  //1 output position
  string verid="mesh-geometry-vertex";
  {
    RapidXMLNode vertexNode=meshNode.AppendNode("vertices");
    vertexNode.SetAttribute("id",ToId(verid));
    string sourceid="mesh-geometry-position";
    vertexNode.AppendNode("input")<<make_pair("semantic","POSITION")<<make_pair("source",ToSource(sourceid));
    
    RapidXMLNode sourceNode=meshNode.AppendNode("source");
    sourceNode.SetAttribute("id",ToId(sourceid));
    RapidXMLNode arrayNode=sourceNode.AppendNode("float_array");
    arrayNode.SetAttribute("id",ToId(sourceid)+"-array");
    arrayNode.SetAttribute("count",mesh.pos_.size()*3);
    ostringstream oss;
    for (zuint j=0; j<mesh.pos_.size(); j++) 
      oss<<mesh.pos_[j];
    arrayNode<<oss.str();

    RapidXMLNode accNode=sourceNode.AppendNode("technique_common").AppendNode("accessor");
    accNode.SetAttribute("source",ToSource(sourceid+"-array"));
    accNode.SetAttribute("count",mesh.pos_.size());
    accNode.SetAttribute("stride","3");
    accNode.AppendNode("param")<<make_pair("name","X")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","Y")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","Z")<<make_pair("type","float");
  }


  //2 output uv
  string uvsourceid="mesh-geometry-uv";
  if (!mesh.tex_.empty())
  {
    RapidXMLNode sourceNode=meshNode.AppendNode("source");
    sourceNode.SetAttribute("id",ToId(uvsourceid));
    RapidXMLNode arrayNode=sourceNode.AppendNode("float_array");
    arrayNode.SetAttribute("id",ToId(uvsourceid)+"-array");
    arrayNode.SetAttribute("count",mesh.tex_.size()*2);
    ostringstream oss;
    for (zuint j=0; j<mesh.tex_.size(); j++) 
      oss<<mesh.tex_[j][0]<<' '<<mesh.tex_[j][1]<<' ';
    arrayNode<<oss.str();

    RapidXMLNode accNode=sourceNode.AppendNode("technique_common").AppendNode("accessor");
    accNode.SetAttribute("source",ToSource(ToId(uvsourceid)+"-array"));
    accNode.SetAttribute("count",mesh.tex_.size());
    accNode.SetAttribute("stride","2");
    accNode.AppendNode("param")<<make_pair("name","S")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","T")<<make_pair("type","float");
  }

  //3 output normal
  string norsourceid="mesh-geometry-normal";
  if (!mesh.nor_.empty())
  {
    RapidXMLNode sourceNode=meshNode.AppendNode("source");
    sourceNode.SetAttribute("id",ToId(norsourceid));
    RapidXMLNode arrayNode=sourceNode.AppendNode("float_array");
    arrayNode.SetAttribute("id",ToId(norsourceid)+"-array");
    arrayNode.SetAttribute("count",mesh.nor_.size()*3);
    ostringstream oss;
    for (zuint j=0; j<mesh.nor_.size(); j++) 
      oss<<mesh.nor_[j];
    arrayNode<<oss.str();

    RapidXMLNode accNode=sourceNode.AppendNode("technique_common").AppendNode("accessor");
    accNode.SetAttribute("source",ToSource(ToId(norsourceid)+"-array"));
    accNode.SetAttribute("count",mesh.nor_.size());
    accNode.SetAttribute("stride","3");
    accNode.AppendNode("param")<<make_pair("name","X")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","Y")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","Z")<<make_pair("type","float");
  }

  //4 output color
  string colorsourceid="mesh-geometry-color";
  {
    RapidXMLNode sourceNode=meshNode.AppendNode("source");
    sourceNode.SetAttribute("id",ToId(colorsourceid));
    RapidXMLNode arrayNode=sourceNode.AppendNode("float_array");
    arrayNode.SetAttribute("id",ToId(colorsourceid)+"-array");
    arrayNode.SetAttribute("count",mesh.color_.size()*3);
    ostringstream oss;
    for (zuint j=0; j<mesh.color_.size(); j++) 
      oss<<mesh.tex_[j];
    arrayNode<<oss.str();

    RapidXMLNode accNode=sourceNode.AppendNode("technique_common").AppendNode("accessor");
    accNode.SetAttribute("source",ToSource(ToId(colorsourceid)+"-array"));
    accNode.SetAttribute("count",mesh.color_.size());
    accNode.SetAttribute("stride","3");
    accNode.AppendNode("param")<<make_pair("name","R")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","G")<<make_pair("type","float");
    accNode.AppendNode("param")<<make_pair("name","B")<<make_pair("type","float");
  }

  //4 output triangles
  for (zuint groupi=0;groupi<mesh.groups_.size();groupi++)
  {
    int inputn=1;
    int inputt[10];
    inputt[0]=SEM_VERTEX;
    TriMesh::MeshGroup *g=mesh.groups_[groupi];
    if (!g->facet_.empty()) inputt[inputn++]=SEM_TEXCOORD;
    if (!g->facen_.empty()) inputt[inputn++]=SEM_NORMAL;
    if (!g->facec_.empty()) inputt[inputn++]=SEM_COLOR;
    RapidXMLNode triNode=meshNode.AppendNode("triangles");
    for (int i=0; i<inputn; i++)
    {
      switch(inputt[i])
      {
      case SEM_VERTEX:
        triNode.AppendNode("input")<<make_pair("semantic","VERTEX")<<make_pair("source",ToSource(verid))<<make_pair("offset",ToString(i));
        break;
      case SEM_TEXCOORD:
        triNode.AppendNode("input")<<make_pair("semantic","TEXCOORD")<<make_pair("source",ToSource(uvsourceid))<<make_pair("offset",ToString(i));
        break;
      case SEM_NORMAL:
        triNode.AppendNode("input")<<make_pair("semantic","NORMAL")<<make_pair("source",ToSource(norsourceid))<<make_pair("offset",ToString(i));
        break;
      case SEM_COLOR:
        triNode.AppendNode("input")<<make_pair("semantic","COLOR")<<make_pair("source",ToSource(colorsourceid))<<make_pair("offset",ToString(i));
        break;
      }
    }
    int trin=mesh.groups_[groupi]->facep_.size();
    if (!g->matname_.empty())
      triNode.SetAttribute("material",ToId(g->matname_));
    triNode.SetAttribute("count",g->facep_.size());
    {
      ostringstream oss;
      for (zuint i=0; i<g->facep_.size(); i++) for (zuint j=0; j<3; j++)
      {
        for (int k=0; k<inputn; k++)
        {
          switch(inputt[k])
          {
          case SEM_VERTEX: oss<<g->facep_[i][j]<<' '; break;
          case SEM_TEXCOORD: oss<<g->facet_[i][j]<<' '; break;
          case SEM_NORMAL: oss<<g->facen_[i][j]<<' '; break;
          case SEM_COLOR: oss<<g->facec_[i][j]<<' '; break;
          }
        }
      }
      triNode.AppendNode("p")<<oss.str();
    }
    if (save_mat_ && !mesh.groups_[groupi]->matname_.empty())
      WriteMaterial(mesh,dae,mesh.groups_[groupi]->matname_);
  }
}


void DaeIO::WriteVisual(TriMesh &mesh, DaeData &dae, string &url)
{
  RapidXMLNode sceneNode=dae.libVisualsNode.AppendNode("visual_scene");
  sceneNode<<make_pair("id",ToId(url));
  sceneNode<<make_pair("name",ToId(url));

  RapidXMLNode modelNode=sceneNode.AppendNode("node");
  modelNode<<make_pair("id","Model0");
  modelNode<<make_pair("name","Model0");

  string geometryUrl="mesh-geometry";
  modelNode.AppendNode("instance_geometry").SetAttribute("url",ToSource(geometryUrl));
  WriteGeometry(mesh,dae,geometryUrl);
}

bool DaeIO::Save(const string &filename, TriMesh &mesh, bool savetex)
{
  save_mat_ = savetex;
  RapidXML xml;
  RapidXMLNode head=xml.AppendNode("COLLADA");
  head<<make_pair("xmlns","http://www.collada.org/2005/11/COLLADASchema");
  head<<make_pair("version","1.4.0");

  //assert
  RapidXMLNode assetnode=head.AppendNode("asset");
  RapidXMLNode contributornode=assetnode.AppendNode("contributor");
  contributornode.AppendNode("authoring_tool")<<"zzzEngine";
  contributornode.AppendNode("author")<<"Zhexi Wang";
  assetnode.AppendNode("up_axis")<<"Z_UP";

  DaeData dae(head.AppendNode("library_images"),
        head.AppendNode("library_materials"),
        head.AppendNode("library_effects"),
        head.AppendNode("library_geometries"),
        head.AppendNode("library_visual_scenes"));
  RapidXMLNode sceneNode=head.AppendNode("scene");
  string visualUrl=GetBase(filename)+"Scene";
  sceneNode.AppendNode("instance_visual_scene")<<make_pair("url",ToSource(visualUrl));

  WriteVisual(mesh,dae,visualUrl);


  xml.SaveFile(filename);
  return true;
}

}
